"""
Side-by-side visualization of raw vs normalized Hi-C matrices from .npz files.

- Loads dense or scipy.sparse NPZ
- Applies chosen normalization (oe/minmax/log/zscore/ice/none)
- Saves a 1x2 figure: left=raw (display-only log1p optional), right=normalized
"""
import os
import argparse
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser(description="Visualize raw vs normalized Hi-C from .npz")
    parser.add_argument("--inputs", type=str, required=True, help="Comma-separated .npz paths or a directory (recursive search)")
    parser.add_argument("--output_dir", type=str, required=True, help="Output directory for PNGs")
    parser.add_argument("--chromosome", type=str, default="", help="Dense NPZ key (e.g., chr1). Empty = first key")
    parser.add_argument("--norm_method", type=str, default="oe", choices=["oe","minmax","log","zscore","ice","none"], help="Normalization method")
    parser.add_argument("--norm_log", type=str, default="True", help="Apply log1p inside normalization (True/False)")
    parser.add_argument("--display_log_raw", type=str, default="True", help="Use log1p for RAW display (True/False)")
    parser.add_argument("--vmin_pct", type=float, default=2.0, help="Lower percentile for display (0-100), -1 disables")
    parser.add_argument("--vmax_pct", type=float, default=98.0, help="Upper percentile for display (0-100), -1 disables")
    parser.add_argument("--window_size", type=int, default=0, help="If >0, crop a window of this size (center by default)")
    parser.add_argument("--start_h", type=int, default=-1, help="Optional top-left row for window")
    parser.add_argument("--start_w", type=int, default=-1, help="Optional top-left col for window")
    parser.add_argument("--limit", type=int, default=8, help="Max number of files to visualize")
    args = parser.parse_args()
    args.norm_log = str(args.norm_log).lower() in ("1","true","t","yes","y")
    args.display_log_raw = str(args.display_log_raw).lower() in ("1","true","t","yes","y")
    return args


def list_npz_files(inputs: str):
    paths = []
    parts = [p.strip() for p in inputs.split(",") if p.strip()]
    for p in parts:
        if os.path.isdir(p):
            for root, _, files in os.walk(p):
                for f in files:
                    if f.endswith(".npz"):
                        paths.append(os.path.join(root, f))
        else:
            if p.endswith(".npz") and os.path.exists(p):
                paths.append(p)
    # de-duplicate while preserving order
    return list(dict.fromkeys(paths))


def load_matrix_from_npz(npz_path: str, chromosome: str = ""):
    data = np.load(npz_path)
    keys = set(list(data.keys()))
    # scipy.sparse NPZ
    if {"data", "indices", "indptr", "shape"}.issubset(keys):
        from scipy import sparse
        return sparse.load_npz(npz_path).toarray()
    # dense NPZ
    if chromosome:
        chr_name = str(chromosome).replace("chr", "")
        candidates = [
            chromosome,
            f"chr{chr_name}",
            chr_name,
            f"matrix_{chromosome}",
            f"matrix_chr{chr_name}",
            f"matrix_{chr_name}",
        ]
        for k in candidates:
            if k in data:
                mat = data[k]
                break
        else:
            mat = data[list(data.keys())[0]]
    else:
        mat = data[list(data.keys())[0]]
    if mat.ndim != 2:
        raise ValueError(f"Matrix must be 2D, got shape {mat.shape} in {npz_path}")
    return mat


def normalize_hic_matrix_np(matrix: np.ndarray, method: str = 'oe', log_transform: bool = True) -> np.ndarray:
    """Torch-free normalization mirroring Squidiff.hic_datasets.normalize_hic_matrix."""
    if method == 'none':
        normalized = matrix.astype(np.float32, copy=True)
    elif method == 'log':
        normalized = np.log1p(matrix)
    elif method == 'minmax':
        x = np.log1p(matrix) if log_transform else matrix.astype(np.float64, copy=True)
        flat = x.flatten()
        pos = flat[flat > 0]
        if pos.size > 0:
            lo = np.percentile(pos, 1.0)
            hi = np.percentile(pos, 99.5)
        else:
            lo = np.percentile(flat, 1.0)
            hi = np.percentile(flat, 99.5)
        if hi <= lo:
            hi = lo + 1e-6
        x = np.clip(x, lo, hi)
        normalized = (x - lo) / (hi - lo)
    elif method == 'zscore':
        x = np.log1p(matrix) if log_transform else matrix.astype(np.float64, copy=True)
        flat = x.flatten()
        pos = flat[flat > 0]
        if pos.size > 1 and pos.std() > 0:
            mu = flat.mean()
            sd = flat.std()
            if sd <= 0:
                normalized = x
            else:
                normalized = (x - mu) / sd
            normalized = np.nan_to_num(normalized, 0)
        else:
            normalized = x
    elif method == 'oe':
        x = matrix.astype(np.float64, copy=True)
        n = x.shape[0]
        dists = np.abs(np.subtract.outer(np.arange(n), np.arange(n)))
        expected = np.zeros(n, dtype=np.float64)
        for d in range(n):
            diag_vals = np.diagonal(x, offset=d)
            nz = diag_vals[diag_vals > 0]
            expected[d] = nz.mean() if nz.size > 0 else 0.0
        expected[expected == 0] = 1.0
        normalized = x / expected[dists]
        normalized = np.nan_to_num(normalized, 0.0, posinf=0.0, neginf=0.0)
        if log_transform:
            normalized = np.log1p(normalized)
        flat = normalized.flatten()
        pos = flat[flat > 0]
        if pos.size > 0:
            lo = np.percentile(pos, 1.0)
            hi = np.percentile(pos, 99.5)
        else:
            lo = np.percentile(flat, 1.0)
            hi = np.percentile(flat, 99.5)
        if hi <= lo:
            hi = lo + 1e-6
        normalized = np.clip(normalized, lo, hi)
        normalized = (normalized - lo) / (hi - lo)
    elif method == 'ice':
        x = np.log1p(matrix) if log_transform else matrix.astype(np.float64, copy=True)
        row_sums = x.sum(axis=1)
        col_sums = x.sum(axis=0)
        row_sums[row_sums == 0] = 1
        col_sums[col_sums == 0] = 1
        normalized = x / np.sqrt(np.outer(row_sums, col_sums))
        normalized = np.nan_to_num(normalized, 0)
    else:
        raise ValueError(f"Unknown normalization method: {method}")
    return normalized.astype(np.float32)


def extract_window(matrix: np.ndarray, ws: int, sh: int = -1, sw: int = -1) -> np.ndarray:
    H, W = matrix.shape
    ws = int(ws)
    if ws <= 0:
        return matrix
    pad_h = max(0, ws - H)
    pad_w = max(0, ws - W)
    if pad_h > 0 or pad_w > 0:
        matrix = np.pad(matrix, ((0, pad_h), (0, pad_w)), mode="constant", constant_values=0)
        H, W = matrix.shape
    if sh < 0 or sw < 0:
        sh = max(0, (H - ws) // 2)
        sw = max(0, (W - ws) // 2)
    else:
        sh = max(0, min(int(sh), max(0, H - ws)))
        sw = max(0, min(int(sw), max(0, W - ws)))
    return matrix[sh:sh + ws, sw:sw + ws]


def compute_display_range(a: np.ndarray, b: np.ndarray, vmin_pct: float, vmax_pct: float):
    if vmin_pct < 0 or vmax_pct < 0:
        return None, None
    cat = np.concatenate([a.flatten(), b.flatten()])
    vmin = float(np.percentile(cat, vmin_pct))
    vmax = float(np.percentile(cat, vmax_pct))
    if vmin >= vmax:
        return None, None
    return vmin, vmax


def visualize_pair(raw_mat: np.ndarray, norm_mat: np.ndarray, save_path: str, vmin_pct: float, vmax_pct: float):
    try:
        import matplotlib.pyplot as plt
        vmin, vmax = compute_display_range(raw_mat, norm_mat, vmin_pct, vmax_pct)
        fig, axes = plt.subplots(1, 2, figsize=(10, 5), dpi=150)
        im0 = axes[0].imshow(raw_mat, cmap="Reds", origin="lower", vmin=vmin, vmax=vmax)
        axes[0].set_title("Raw (display)")
        axes[0].axis("off")
        im1 = axes[1].imshow(norm_mat, cmap="Reds", origin="lower", vmin=vmin, vmax=vmax)
        axes[1].set_title("Normalized")
        axes[1].axis("off")
        fig.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04)
        fig.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        fig.tight_layout()
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
    except Exception as e:
        print(f"[WARN] Failed to plot {save_path}: {e}")


def main():
    args = parse_args()
    files = list_npz_files(args.inputs)
    os.makedirs(args.output_dir, exist_ok=True)
    count = 0
    for fpath in files:
        if count >= args.limit:
            break
        try:
            mat = load_matrix_from_npz(fpath, args.chromosome)
        except Exception as e:
            print(f"[SKIP] {fpath}: {e}")
            continue

        # Optional window crop for fair comparison
        if args.window_size and int(args.window_size) > 0:
            mat = extract_window(mat, int(args.window_size), int(args.start_h), int(args.start_w))

        # Prepare raw display matrix (display-only transform)
        raw_disp = np.log1p(mat) if args.display_log_raw else mat.astype(np.float32)

        # Normalize via dataset function
        norm = normalize_hic_matrix_np(mat, method=args.norm_method, log_transform=args.norm_log)

        # Save figure
        base = os.path.splitext(os.path.basename(fpath))[0]
        tag = f"_{args.norm_method}{'_log' if args.norm_log else ''}"
        if args.window_size and int(args.window_size) > 0:
            tag += f"_ws{int(args.window_size)}"
        save_path = os.path.join(args.output_dir, f"{base}_raw_vs_norm{tag}.png")
        visualize_pair(raw_disp, norm, save_path, args.vmin_pct, args.vmax_pct)
        print(f"[OK] {save_path}")
        count += 1

    print(f"Saved {count} figures -> {os.path.abspath(args.output_dir)}")


if __name__ == "__main__":
    main()


